from torch.utils.data import Dataset



class EvalSubset(Dataset):
    def __init__(self, dataset, indices=None):
        self.dataset = dataset
        if indices is None:
            self.indices = list(range(len(dataset)))
        else:
            self.indices = indices

    def __getitem__(self, index):
        actual_index = self.indices[index]
        img, label = self.dataset[actual_index]
        return img, label,

    def __len__(self):
        return len(self.indices)

    def addEvalRecord(self, judge_result):
        self.dataset.addNewEvalRecord(self.indices, judge_result)

    def addConfidenceRecord(self, confidences_result, beta):
        self.dataset.addConfidencesRecord(self.indices, confidences_result, beta)

    def updateIndices(self, update_type, sample_ratio=None):
        """

        :param update_type:
        :param sample_ratio: 表示从easy examples中抽样出的比例
        :return:
        """
        self.indices = self.dataset.getIndicesFromEvalDataset(update_type, sample_ratio)